package com.ruoyi.framework.encrypt.utils;

import org.apache.commons.codec.binary.Base64;
import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom;
import java.util.HashMap;
import java.util.Map;

/**
 * AES 加密工具类
 * 模式：AES/CBC/NoPadding + 填充字符强制过滤
 */
public class AESUtils {
    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/CBC/NoPadding";
    private static final int KEY_SIZE = 128;
    private static final int IV_LENGTH = 16;
    private static final int BLOCK_SIZE = 16;

    public static Map<String, String> generateAesKeyAndIv() throws Exception {
        KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM);
        SecureRandom secureRandom = new SecureRandom();
        keyGenerator.init(KEY_SIZE, secureRandom);
        SecretKey secretKey = keyGenerator.generateKey();
        String aesKeyBase64 = Base64.encodeBase64String(secretKey.getEncoded());

        byte[] iv = new byte[IV_LENGTH];
        secureRandom.nextBytes(iv);
        String aesIvBase64 = Base64.encodeBase64String(iv);

        Map<String, String> result = new HashMap<>(2);
        result.put("AES_KEY", aesKeyBase64);
        result.put("AES_IV", aesIvBase64);
        return result;
    }

    public static String aesEncrypt(String plainText, String aesKeyBase64, String aesIvBase64) throws Exception {
        byte[] keyBytes = Base64.decodeBase64(aesKeyBase64);
        byte[] ivBytes = Base64.decodeBase64(aesIvBase64);

        if (keyBytes.length != KEY_SIZE / 8) throw new IllegalArgumentException("AES密钥必须16字节");
        if (ivBytes.length != IV_LENGTH) throw new IllegalArgumentException("IV必须16字节");

        // 手动PKCS5填充（与JS端一致）
        byte[] plainBytes = plainText.getBytes("UTF-8");
        int paddingLen = BLOCK_SIZE - (plainBytes.length % BLOCK_SIZE);
        byte[] paddedBytes = new byte[plainBytes.length + paddingLen];
        System.arraycopy(plainBytes, 0, paddedBytes, 0, plainBytes.length);
        for (int i = plainBytes.length; i < paddedBytes.length; i++) {
            paddedBytes[i] = (byte) paddingLen;
        }

        SecretKeySpec keySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        IvParameterSpec ivSpec = new IvParameterSpec(ivBytes);
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);

        byte[] encryptBytes = cipher.doFinal(paddedBytes);
        return Base64.encodeBase64String(encryptBytes);
    }

    /**
     * 解密核心：先去填充，再强制过滤ASCII 1~16的填充字符
     */
    public static String aesDecrypt(String cipherTextBase64, String aesKeyBase64, String aesIvBase64) throws Exception {
        byte[] keyBytes = Base64.decodeBase64(aesKeyBase64);
        byte[] ivBytes = Base64.decodeBase64(aesIvBase64);
        byte[] cipherBytes = Base64.decodeBase64(cipherTextBase64);

        if (keyBytes.length != KEY_SIZE / 8) throw new IllegalArgumentException("AES密钥必须16字节");
        if (ivBytes.length != IV_LENGTH) throw new IllegalArgumentException("IV必须16字节");
        if (cipherBytes.length % BLOCK_SIZE != 0) throw new IllegalArgumentException("密文必须是16的整数倍");

        // 1. AES解密
        SecretKeySpec keySpec = new SecretKeySpec(keyBytes, ALGORITHM);
        IvParameterSpec ivSpec = new IvParameterSpec(ivBytes);
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);
        byte[] decryptBytes = cipher.doFinal(cipherBytes);

        // 2. 尝试手动去填充（优先按标准处理）
        byte[] plainBytes;
        try {
            plainBytes = removePkcs5Padding(decryptBytes);
        } catch (Exception e) {
            // 填充校验失败时，直接使用解密后的字节（后续强制过滤）
            plainBytes = decryptBytes;
        }

        // 3. 强制过滤ASCII 1~16的不可见字符（PKCS5填充字符）
        byte[] filteredBytes = filterPaddingChars(plainBytes);

        // 4. UTF-8解码（还原汉字）
        return new String(filteredBytes, "UTF-8");
    }

    /**
     * 标准PKCS5去填充
     */
    private static byte[] removePkcs5Padding(byte[] data) throws Exception {
        if (data == null || data.length == 0) return data;
        int paddingLen = data[data.length - 1] & 0xFF;
        if (paddingLen < 1 || paddingLen > BLOCK_SIZE) throw new Exception("填充长度异常");
        for (int i = data.length - paddingLen; i < data.length; i++) {
            if ((data[i] & 0xFF) != paddingLen) throw new Exception("填充格式错误");
        }
        byte[] plainData = new byte[data.length - paddingLen];
        System.arraycopy(data, 0, plainData, 0, plainData.length);
        return plainData;
    }

    /**
     * 强制过滤PKCS5填充字符（ASCII 1~16）
     * 核心：保留所有有效字符（汉字、字母、数字、符号），仅删除填充残留
     */
    private static byte[] filterPaddingChars(byte[] data) {
        if (data == null || data.length == 0) return data;
        int validLen = 0;
        // 先统计有效字节数（排除ASCII 1~16）
        for (byte b : data) {
            int ascii = b & 0xFF;
            if (ascii < 1 || ascii > 16) { // 仅保留非填充字符
                validLen++;
            }
        }
        // 构建过滤后的字节数组
        byte[] filtered = new byte[validLen];
        int idx = 0;
        for (byte b : data) {
            int ascii = b & 0xFF;
            if (ascii < 1 || ascii > 16) {
                filtered[idx++] = b;
            }
        }
        return filtered;
    }

}